module Zeno.Evaluation ( applyTermDefinitions,
                         applyFunction,
                         strictVar,
                         strictVars ) where

import Prelude ()
import StdImports
import qualified Data.Set as Set
                         
import Zeno.Core
import Zeno.Clause
import Zeno.Term
import Zeno.Function

strictVars :: WithinTraversable ZTerm t => t -> [ZVar]
strictVars = nub . mapMaybe strictVar . withinList
          
strictVar :: ZTerm -> Maybe ZVar
strictVar var@(Var zvar)
  | isVariableClass (zVarClass zvar) 
      && destructibleVar zvar = Just zvar
  | otherwise = Nothing
strictVar term =
  case applyFunction [] term of
    Nothing -> Nothing
    Just (Value term') -> strictVar term'
    Just (CaseOf _ term' _) -> strictVar term'
                         
applyTermDefinitions :: WithinTraversable ZTerm (t ZVar) => 
    [ZEquality] -> t ZVar -> Writer Any (t ZVar)
applyTermDefinitions defs t = do
  (t', any) <- listen 
    $ mapWithinM applyTermDefinition t
  if getAny any 
    then applyTermDefinitions defs t'
    else return t'
  where
    applyTermDefinition :: ZTerm -> Writer Any ZTerm
    applyTermDefinition term =
      case applyFunction defs term of 
        Just (Value term') -> tell (Any True) >> return term'
        _ -> return term

-- This function may give an infinite loop for 
--   non-primitive recursive function definitions
applyFunction :: [ZEquality] -> ZTerm -> Maybe ZExpression
applyFunction defs term =
  case (zVarClass func_var, null args) of
    (ZFunction func, False) ->
      let params = map Var $ funcParameters func
      in if length params /= length args
        then Nothing
        else
          let arg_mapping = Set.fromList $ params `zip` args
              expr = funcExpression func
              mapped_expr = substituteWithin arg_mapping expr
          in Just $ simplifyExpression mapped_expr
    _ -> Nothing
  where
    (Var func_var) : args = flattenTerm term
  
    simplifyExpression :: ZExpression -> ZExpression
    simplifyExpression expr@(Value term) = expr              
    simplifyExpression expr@(CaseOf id term whens) =
      -- Multiple matches here indicate a contradiction in the goal antecedents
      let simpl = case mapMaybe matchWhen whens of
            [expr'] -> simplifyExpression expr'
            [] ->
              let maybe_expr = applyFunction defs term
              in case maybe_expr of
                Just (Value term') ->
                  let expr' = CaseOf id term' whens
                  in if constructorTerm term 
                    then simplifyExpression expr'
                    else expr'
                _ -> expr
            _ -> expr
          
          expr' = substituteWithin (equalitiesToSubstitution defs) simpl
      in if expr' /= expr
        then simplifyExpression expr'
        else expr
      where
        ((Var func) : args) = flattenTerm term
          
        matchWhen :: ZWhen -> Maybe ZExpression
        matchWhen (When term' expr') = do
          let ((Var func') : args') = flattenTerm term'
          guard (func' == func) 
          let arg_mapping = Set.fromList $ args' `zip` args
          return $ substituteWithin arg_mapping expr'
